import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split, GridSearchCV, cross_val_score
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression, Ridge, Lasso
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from sklearn.feature_selection import SelectKBest, f_regression
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

class RevenuePredictionModel:
    def __init__(self, file_path):
        self.file_path = file_path
        self.df = None
        self.X_train = None
        self.X_test = None
        self.y_train = None
        self.y_test = None
        self.scaler = StandardScaler()
        self.models = {}
        self.best_model = None
        self.best_model_name = None
        self.feature_importance = None
    
    def load_data(self):
        """加载和预处理数据"""
        print("正在加载数据...")
        self.df = pd.read_excel(self.file_path)
        
        # 数据基本信息
        print(f"数据形状: {self.df.shape}")
        
        # 特征工程 - 创建新的财务指标
        print("正在进行特征工程...")
        
        # 毛利率
        self.df['毛利率'] = (self.df['营业收入'] - self.df['营业成本']) / self.df['营业收入']
        
        # 资产负债率
        self.df['资产负债率'] = self.df['负债合计'] / self.df['资产总计']
        
        # 所有者权益比率
        self.df['所有者权益比率'] = self.df['所有者权益合计'] / self.df['资产总计']
        
        # 现金流量比率
        self.df['经营现金流比率'] = self.df['经营活动产生的现金流量净额'] / self.df['营业收入']
        
        # 现金净增加额比率
        self.df['现金净增加额比率'] = self.df['现金及现金等价物净增加额'] / self.df['营业收入']
        
        # 投资活动现金流比率
        self.df['投资现金流比率'] = self.df['投资活动产生的现金流量净额'] / self.df['营业收入']
        
        # 筹资活动现金流比率
        self.df['筹资现金流比率'] = self.df['筹资活动产生的现金流量净额'] / self.df['营业收入']
        
        # 处理可能的无穷大和NaN值
        self.df.replace([np.inf, -np.inf], np.nan, inplace=True)
        self.df.fillna(0, inplace=True)
        
        print("特征工程完成！新增7个财务指标")
    
    def split_data(self):
        """划分训练集和测试集"""
        # 选择特征和目标变量
        feature_columns = [
            '营业成本', '资产总计', '负债合计', '所有者权益合计',
            '经营活动产生的现金流量净额', '投资活动产生的现金流量净额',
            '筹资活动产生的现金流量净额', '现金及现金等价物净增加额',
            '毛利率', '资产负债率', '所有者权益比率', '经营现金流比率',
            '现金净增加额比率', '投资现金流比率', '筹资现金流比率'
        ]
        
        X = self.df[feature_columns]
        y = self.df['营业收入']
        
        # 划分训练集和测试集
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(
            X, y, test_size=0.3, random_state=42
        )
        
        # 数据标准化
        self.X_train = self.scaler.fit_transform(self.X_train)
        self.X_test = self.scaler.transform(self.X_test)
        
        print(f"训练集大小: {self.X_train.shape[0]}")
        print(f"测试集大小: {self.X_test.shape[0]}")
    
    def train_models(self):
        """训练多种预测模型"""
        print("正在训练多种预测模型...")
        
        # 线性回归
        lr = LinearRegression()
        lr.fit(self.X_train, self.y_train)
        self.models['线性回归'] = lr
        
        # Ridge回归
        ridge = Ridge(alpha=1.0)
        ridge.fit(self.X_train, self.y_train)
        self.models['Ridge回归'] = ridge
        
        # Lasso回归
        lasso = Lasso(alpha=0.1)
        lasso.fit(self.X_train, self.y_train)
        self.models['Lasso回归'] = lasso
        
        # 随机森林回归
        rf = RandomForestRegressor(n_estimators=100, random_state=42)
        rf.fit(self.X_train, self.y_train)
        self.models['随机森林'] = rf
        
        # 梯度提升回归
        gbr = GradientBoostingRegressor(n_estimators=100, random_state=42)
        gbr.fit(self.X_train, self.y_train)
        self.models['梯度提升'] = gbr
        
        print("所有模型训练完成！")
    
    def evaluate_models(self):
        """评估模型性能并选择最佳模型"""
        print("\n模型性能评估:")
        print("-" * 80)
        print(f"{'模型':<10} {'R²分数':<10} {'MSE':<20} {'RMSE':<20} {'MAE':<20}")
        print("-" * 80)
        
        best_r2 = -float('inf')
        
        for name, model in self.models.items():
            y_pred = model.predict(self.X_test)
            r2 = r2_score(self.y_test, y_pred)
            mse = mean_squared_error(self.y_test, y_pred)
            rmse = np.sqrt(mse)
            mae = mean_absolute_error(self.y_test, y_pred)
            
            # 交叉验证
            cv_scores = cross_val_score(model, self.X_train, self.y_train, cv=5, scoring='r2')
            cv_mean = cv_scores.mean()
            
            print(f"{name:<10} {r2:.4f}      {mse:.2e}     {rmse:.2e}     {mae:.2e}     CV-R²: {cv_mean:.4f}")
            
            # 更新最佳模型
            if r2 > best_r2:
                best_r2 = r2
                self.best_model = model
                self.best_model_name = name
        
        print("-" * 80)
        print(f"最佳模型: {self.best_model_name} (R² = {best_r2:.4f})")
    
    def feature_analysis(self):
        """分析特征重要性"""
        print("\n特征重要性分析:")
        
        feature_columns = [
            '营业成本', '资产总计', '负债合计', '所有者权益合计',
            '经营活动产生的现金流量净额', '投资活动产生的现金流量净额',
            '筹资活动产生的现金流量净额', '现金及现金等价物净增加额',
            '毛利率', '资产负债率', '所有者权益比率', '经营现金流比率',
            '现金净增加额比率', '投资现金流比率', '筹资现金流比率'
        ]
        
        # 根据不同模型类型获取特征重要性
        if self.best_model_name in ['线性回归', 'Ridge回归', 'Lasso回归']:
            importance = np.abs(self.best_model.coef_)
        else:  # 树模型
            importance = self.best_model.feature_importances_
        
        # 创建特征重要性字典
        self.feature_importance = dict(zip(feature_columns, importance))
        
        # 排序并打印
        sorted_importance = sorted(self.feature_importance.items(), key=lambda x: x[1], reverse=True)
        
        print(f"最佳模型 {self.best_model_name} 的特征重要性:")
        for i, (feature, importance) in enumerate(sorted_importance, 1):
            print(f"{i}. {feature}: {importance:.4f}")
        
        # 可视化前10个重要特征
        top_features = sorted_importance[:10]
        features, scores = zip(*top_features)
        
        plt.figure(figsize=(12, 8))
        plt.barh(range(len(features)), scores)
        plt.yticks(range(len(features)), features)
        plt.xlabel('重要性得分')
        plt.ylabel('特征')
        plt.title(f'{self.best_model_name} - 前10个重要特征')
        plt.gca().invert_yaxis()
        plt.tight_layout()
        plt.savefig('feature_importance.png', dpi=300)
        print("特征重要性图已保存为 'feature_importance.png'")
    
    def correlation_analysis(self):
        """相关性分析"""
        print("\n相关性分析:")
        
        # 选择关键特征进行相关性分析
        key_features = [
            '营业收入', '营业成本', '资产总计', '负债合计', '所有者权益合计',
            '毛利率', '资产负债率', '经营现金流比率'
        ]
        
        corr_matrix = self.df[key_features].corr()
        
        # 可视化相关性热图
        plt.figure(figsize=(12, 10))
        sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', fmt='.2f', square=True)
        plt.title('关键特征相关性热图')
        plt.tight_layout()
        plt.savefig('correlation_heatmap.png', dpi=300)
        print("相关性热图已保存为 'correlation_heatmap.png'")
        
        # 打印与营业收入的相关性
        revenue_corr = corr_matrix['营业收入'].sort_values(ascending=False)
        print("与营业收入的相关性:")
        for feature, corr in revenue_corr.items():
            if feature != '营业收入':
                print(f"{feature}: {corr:.4f}")
    
    def model_diagnostics(self):
        """模型诊断"""
        print("\n模型诊断:")
        
        # 获取最佳模型的预测结果
        y_pred = self.best_model.predict(self.X_test)
        
        # 残差分析
        residuals = self.y_test - y_pred
        
        # 残差分布图
        plt.figure(figsize=(12, 6))
        plt.subplot(1, 2, 1)
        sns.distplot(residuals, kde=True)
        plt.xlabel('残差')
        plt.ylabel('频数')
        plt.title('残差分布')
        
        # 预测值vs实际值散点图
        plt.subplot(1, 2, 2)
        plt.scatter(self.y_test, y_pred, alpha=0.5)
        max_val = max(max(self.y_test), max(y_pred))
        plt.plot([0, max_val], [0, max_val], 'r--')
        plt.xlabel('实际值')
        plt.ylabel('预测值')
        plt.title('预测值 vs 实际值')
        
        plt.tight_layout()
        plt.savefig('model_diagnostics.png', dpi=300)
        print("模型诊断图已保存为 'model_diagnostics.png'")
    
    def hyperparameter_tuning(self):
        """对最佳模型进行超参数调优"""
        if self.best_model_name == '随机森林':
            print("\n对随机森林模型进行超参数调优...")
            param_grid = {
                'n_estimators': [50, 100, 200],
                'max_depth': [None, 10, 20, 30],
                'min_samples_split': [2, 5, 10],
                'min_samples_leaf': [1, 2, 4]
            }
            
            grid_search = GridSearchCV(
                estimator=RandomForestRegressor(random_state=42),
                param_grid=param_grid,
                cv=5,
                scoring='r2',
                n_jobs=-1
            )
            
            grid_search.fit(self.X_train, self.y_train)
            
            print(f"最佳参数: {grid_search.best_params_}")
            print(f"调优后的R²分数: {grid_search.best_score_:.4f}")
            
            # 更新最佳模型
            self.best_model = grid_search.best_estimator_
    
    def run_pipeline(self):
        """运行完整的预测分析流程"""
        print("===== 上市公司营收预测模型分析 ====")
        
        # 1. 加载数据
        self.load_data()
        
        # 2. 划分数据
        self.split_data()
        
        # 3. 训练多个模型
        self.train_models()
        
        # 4. 评估模型
        self.evaluate_models()
        
        # 5. 超参数调优
        self.hyperparameter_tuning()
        
        # 6. 特征重要性分析
        self.feature_analysis()
        
        # 7. 相关性分析
        self.correlation_analysis()
        
        # 8. 模型诊断
        self.model_diagnostics()
        
        print("\n===== 分析完成 ====")
        print(f"最佳预测模型: {self.best_model_name}")
        
        # 生成关键发现
        self.generate_key_findings()
    
    def generate_key_findings(self):
        """生成关键发现"""
        print("\n===== 关键发现 ====")
        
        # 1. 预测模型性能
        y_pred = self.best_model.predict(self.X_test)
        r2 = r2_score(self.y_test, y_pred)
        print(f"1. 最终模型 ({self.best_model_name}) 的R²分数为: {r2:.4f}")
        
        # 2. 最重要的特征
        top_features = sorted(self.feature_importance.items(), key=lambda x: x[1], reverse=True)[:3]
        print("2. 影响营收的最重要三个因素:")
        for i, (feature, importance) in enumerate(top_features, 1):
            print(f"   {i}. {feature} (重要性得分: {importance:.4f})")
        
        # 3. 相关性洞察
        key_features = ['营业收入', '营业成本', '资产总计', '负债合计', '毛利率']
        corr_matrix = self.df[key_features].corr()
        revenue_corr = corr_matrix['营业收入'].sort_values(ascending=False)
        
        print("3. 与营业收入相关性最高的三个特征:")
        # 转换为列表后切片
        corr_items = list(revenue_corr.items())
        for i, (feature, corr) in enumerate(corr_items[1:4], 1):  # 排除自己
            print(f"   {i}. {feature} (相关系数: {corr:.4f})")
        
        print("\n分析结果已保存为图表文件，可在当前目录查看。")

if __name__ == "__main__":
    # 数据文件路径
    file_path = "../chapter5/上市公司营收数据.xlsx"
    
    # 创建并运行模型
    model = RevenuePredictionModel(file_path)
    model.run_pipeline()